import os, sys, math, shutil

# NOTE : should be relative path.
data_set_root_record = '/home/leo/Downloads/datas/road_car_split1_record'


all_files = []
for file in os.listdir(data_set_root_record):
    if ".record" in os.path.splitext(file)[1]:
        all_files.append(os.path.join(data_set_root_record, file))

eval_idx = len(all_files) - math.ceil(len(all_files) / 10)
# eval_idx = len(all_files)
train_files = all_files[:eval_idx]
eval_files = all_files[eval_idx:]

idx = 0
length = len(train_files)
with open(os.path.join(data_set_root_record,"map.txt"),"w+") as fp:
    for file in train_files:
        dst = os.path.join(data_set_root_record, "road_car_train.record-" + "{:0>5d}-of-{:0>5d}".format(idx, length))
        shutil.move(file, dst)
        idx += 1
        fp.write(os.path.basename(file)+":"+os.path.basename(dst)+"\n")
    idx = 0
    length = len(eval_files)
    for file in eval_files:
        dst = os.path.join(data_set_root_record, "road_car_eval.record-" + "{:0>5d}-of-{:0>5d}".format(idx, length))
        shutil.move(file, dst)
        idx += 1
        fp.write(os.path.basename(file) + ":" + os.path.basename(dst)+"\n")
